
#include "parsebundler.h"

#include <fstream>
#include <iostream>
#include <numeric>
#include <sstream>
#include <string.h>
#include <utility>

#include "config.h"
#include "timer.h"

#include "glog/logging.h"

//using namespace std;
using std::string;
using std::ifstream;
using std::ofstream;
using std::istringstream;
using std::cout;
using std::endl;

// clear all the data
void BundlerParser::ClearData() {
  mFeature_infos.clear();
  mPictures.ClearPicture();
  mCameras.clear();
  bundler_out.clear();
  key_path.clear();
  list_file.clear();
  image_path.clear();
  mMap2dTo3d.clear();
  mPoints3dCost.clear();
  mPoints3D.clear();
}

BundlerParser::BundlerParser() {
  ClearData();
}

BundlerParser::BundlerParser(const std::string& _bundler_out,
                             const std::string& _key_path,
                             const std::string& _list_file,
                             const std::string& _feat_3d_info_file,
                             const std::string& _image_path) {
  bundler_out = _bundler_out;
  key_path = _key_path;
  list_file = _list_file;
  feat_3d_info_file = _feat_3d_info_file;
  image_path = _image_path;
}

BundlerParser::~BundlerParser() {
  ClearData();
}

// parse the bundler file
/*Each camera entry contains the estimated camera intrinsics and extrinsics,
  and has the form:
<f> <k1> <k2>   [the focal length, followed by two radial distortion coeffs]
<R>             [a 3x3 matrix representing the camera rotation]
<t>             [a 3-vector describing the camera translation]

Each point entry has the form
<position>      [a 3-vector describing the 3D position of the point]
<color>         [a 3-vector describing the RGB color of the point]
<view list>     [a list of views the point is visible in]

The view list begins with the length of the list 
(i.e., the number of cameras the point is visible in).
The list is then given as a list of quadruplets <camera> <key> <x> <y>
The pixel positions are floating point numbers in a coordinate system 
where the origin is the top-left corner of the image.
*/
bool BundlerParser::ParseBundlerFile(const std::string& bundler_out, bool bLoad_feat_info) {
  mCameras.clear();
  Timer timer;
  std::ifstream ifs(bundler_out, std::ios::binary | std::ios::_Nocreate);
  if (!ifs.is_open()) {
    LOG(ERROR) << "open bundler fail: " + bundler_out;
    return false;
  }
  std::string line;
  std::getline(ifs, line);  // header
  int num_cam = 0, num_points = 0;
  ifs >> num_cam >> num_points;
  mCameras.resize(num_cam);
  VLOG(1) << "bundler has " << num_cam << " cameras and " << num_points << " 3d points";

  const cv::Matx33d D { 1., 0., 0., 0., -1., 0., 0., 0., -1. };
  for (size_t i = 0; i < num_cam; ++i) {
    ifs >> mCameras[i].f >> mCameras[i].k[0] >> mCameras[i].k[1];
    auto& r = mCameras[i].r;
    ifs >> r(0, 0) >> r(0, 1) >> r(0, 2)
      >> r(1, 0) >> r(1, 1) >> r(1, 2)
      >> r(2, 0) >> r(2, 1) >> r(2, 2);
    auto& t = mCameras[i].t;
    ifs >> t[0] >> t[1] >> t[2];
    // z axis of bundler is opposite to the usual z axis 
    r = D * r;
    t = D * t;
  }
  
  if (!bLoad_feat_info) {
    LOG(INFO) << "load camera pose time: " << timer;
    return true;
  }

  VLOG(2) << "start to load 3d points";
  mFeature_infos.clear();
  mFeature_infos.resize(num_points);

	// load the points, remove feature points in query images 
  int r, g, b, view_lenth;
  for (int i = 0; i < num_points; ++i) {
    ifs >> mFeature_infos[i].mPoint.x
			>> mFeature_infos[i].mPoint.y
			>> mFeature_infos[i].mPoint.z
			>> r >> g >> b;
    mFeature_infos[i].mPoint.r = (unsigned char)r;
    mFeature_infos[i].mPoint.g = (unsigned char)g;
    mFeature_infos[i].mPoint.b = (unsigned char)b;

    ifs >> view_lenth;
		mFeature_infos[i].mView_list.reserve(view_lenth);
    mFeature_infos[i].mDescriptors.reserve(view_lenth);
		for (int j = 0; j < view_lenth; j++) {
      View view;
      ifs >> view.camera >> view.key >> view.keypt.x >> view.keypt.y;
      auto& picture = mPictures.Pictures()[view.camera];
#if   TEST_PNP
      picture.mPoints[view.key] = cv::Point3d{mFeature_infos[i].mPoint.x, mFeature_infos[i].mPoint.y, mFeature_infos[i].mPoint.z};
#endif //  TEST_PNP
      // skip query images
      if (!picture.Flag()) {
          continue;
      }
      // add view and descriptor
      mFeature_infos[i].mView_list.push_back(view);
      auto& desc = picture.Descriptors()[view.key];
      mFeature_infos[i].mDescriptors.push_back(desc);
		}
    // clear invalid 3d points
    if (mFeature_infos[i].mView_list.size() <= 1) {
      mFeature_infos[i].mView_list.clear();
      mFeature_infos[i].mDescriptors.clear();
    }
	}
	ifs.close();
  LOG(INFO) << "load camera pose and points info time: " << timer;
  
	return true;
}


//load the .key info, needed descriptors 
bool BundlerParser::LoadFeatDesc(const std::string& key_path,
                                 const std::string& list_file,
                                 const std::string& image_path,
                                 const bool load_only_query) {
  mPictures.SetPath(key_path, list_file, image_path);
  mPictures.LoadPictures(load_only_query);
  return true;
}

bool BundlerParser::LoadData() {
  Timer timer;
  LOG(INFO) << "try to load from feat_3d_info_file file";
  if (TEST_PNP || (!ParseBundlerFile(bundler_out, false) || !LoadFeatures3DInfoBin(feat_3d_info_file))) {
    LOG(INFO) << "load from feat_3d_info_file fail, load from .feat and .desc file";
    LoadFeatDesc(key_path, list_file, image_path, false);
    ParseBundlerFile(bundler_out, true);

    // then save the feat 3d info
    SaveFeatures3DInfoBin(feat_3d_info_file);
  } else {
    // load only query picture
    LoadFeatDesc(key_path, list_file, image_path, true);
  }
  const auto& parameters = mPictures.Parameters();
  CHECK_EQ(parameters.size(), mCameras.size() * 3) << "camera and parameter size error";
  // load ppx ppy f parameters into Camera
  for (int i = 0; i < mCameras.size(); ++i) {
    mCameras[i].ppx = parameters[3 * i];
    mCameras[i].ppy = parameters[3 * i + 1];
    mCameras[i].f = parameters[3 * i + 2];
  }

  auto lambda = [] (size_t cur_total, const Feature3DInfo& info) {
    return cur_total + info.mView_list.size();
  };
  size_t desc_size = accumulate(mFeature_infos.cbegin(), mFeature_infos.cend(), size_t(0), lambda);
  LOG(INFO) << "load total 2d feature number: " << desc_size;

  // extract desc and 3d points into mDescriptorMat and mPoints3D
  size_t cnt_2d = 0, cnt_3d = 0;
  mPoints3D.clear();
  mPoints3D.reserve(desc_size);
  mPoints3dCost.reserve(desc_size);
  mMap2dTo3d.resize(desc_size);
  mDescriptorMat = cv::Mat(desc_size, Descriptor::desc_legth, CV_8UC1);
  uchar* ptr_data = mDescriptorMat.data;
  for (const auto& feat_info : mFeature_infos) {
    mPoints3dCost.push_back(feat_info.mView_list.size());
    // extract 3d points
    mPoints3D.emplace_back(feat_info.mPoint.x, feat_info.mPoint.y, feat_info.mPoint.z);
    // extract desc mat
    for (const auto& desc : feat_info.mDescriptors) {
      ptr_data = desc.CopyDescData(ptr_data);
      mMap2dTo3d[cnt_2d++] = cnt_3d;
    }
    ++cnt_3d;
  }
  mPoints3D.shrink_to_fit();
  mPoints3dCost.shrink_to_fit();

  // then clear the feature info
  mFeature_infos.clear();

  LOG(INFO) << "load model okay:\n"
    << "camera num: " << mCameras.size() << "\n"
    << "picture num: " << mPictures.PictureNum() << "\n"
    << "3D point num: " << mPoints3D.size();
  LOG(INFO) << "load data time " << timer;
  return true;
}


//save the built information so next time directly load the file
//format::
//#3d points
//for each 3d points, contain the full information
//3dpoint(x, y, z, r g b)
//#view each view(camera key x, y, scale, orientation)
//also all descriptor of one 3d points 
bool BundlerParser::SaveFeatures3DInfo(const std::string& feat_3d_info_file) const {
  CHECK_EQ(mCameras.size(), mPictures.PictureNum()) << "camera and picture number error";
  std::ofstream ofs(feat_3d_info_file, std::ios::trunc | std::ios::binary);
	if (!ofs.is_open()){
    LOG(ERROR) << "open feat_3d_info_file fail: " + feat_3d_info_file;
		return false;
	}
  
  size_t num_points = mFeature_infos.size();
  ofs << num_points << "\n";
	//each 3d points a line
	for (const auto& feat_3d_info : mFeature_infos) {
    ofs << feat_3d_info.mPoint.x << " " << feat_3d_info.mPoint.y << " " << feat_3d_info.mPoint.z << "\n";
		//ofs << feat_3d_info.mPoint.r << " " << feat_3d_info.mPoint.g << " " << feat_3d_info.mPoint.b << "";
		ofs << feat_3d_info.mView_list.size() << " ";
		//save view
		for (const auto& view : feat_3d_info.mView_list) {
      ofs << view.camera << " " << view.key << " " << view.keypt.x << " "
				<< view.keypt.y << " " << view.keypt.scale << " " << view.keypt.orientation << " ";
		}
		ofs << std::endl;

		//save decriptor num_desc = feat_3d_info.mDescriptor.size();
		for (const auto& desc : feat_3d_info.mDescriptors) {
			for (int i = 0; i < desc.desc_legth; ++i){
				ofs << int(desc.ptrDesc[i]) << " ";
			}
      ofs << std::endl;
		}	
	}
	ofs.close();
	return true;
}

bool BundlerParser::LoadFeatures3DInfo(const std::string& feat_3d_info_file) {
  std::ifstream ifs(feat_3d_info_file, std::ios::_Nocreate);
	if (!ifs.is_open()) {
    LOG(INFO) << feat_3d_info_file + " does not exist then reload" ;
		return false;
	}

  size_t num_points = 0;
  ifs >> num_points;
	VLOG(1) << "load feat_3d_info, total num of points:" << num_points << endl;
	
	mFeature_infos.clear();
  mFeature_infos.resize(num_points);
	size_t view_list_len = 0;

	//each 3d points a line
  for (size_t i = 0; i < num_points; ++i) {
		auto& feat_3d_point = mFeature_infos[i].mPoint;
		ifs >> feat_3d_point.x >> feat_3d_point.y >> feat_3d_point.z;
		//is >> feat_3d_point.r >> feat_3d_point.g >> feat_3d_point.b;
		ifs >> view_list_len;

		//load view_list
		mFeature_infos[i].mView_list.resize(view_list_len);
		for (auto& view : mFeature_infos[i].mView_list){
      ifs >> view.camera >> view.key >> view.keypt.x >> view.keypt.y
        >> view.keypt.scale >> view.keypt.orientation;
		}

		//load descriptor num_desc = feat_3d_info.mDescriptor.size();
		mFeature_infos[i].mDescriptors.resize(view_list_len);
		for (int j = 0; j < view_list_len; j++) {
			auto& desc = mFeature_infos[i].mDescriptors[j];
      desc.ptrDesc = new unsigned char[desc.desc_legth];
      if (!desc.ptrDesc) {
				LOG(WARNING) << "new unsigned char fail.try again";
				--j;
				continue;
			}
      int tmp = 0;
      for (int k = 0; k < desc.desc_legth; ++k) {
				ifs >> tmp; 
				desc.ptrDesc[k] = (unsigned char)tmp;
			}
		}
    VLOG(3) << "load the " << i << " 3d point okay";
	}

	return true;
}

bool BundlerParser::SaveFeatures3DInfoBin(const std::string& feat_3d_info_file) const {
  CHECK_EQ(mCameras.size(), mPictures.PictureNum()) << "camera and picture number error";
  std::ofstream ofs(feat_3d_info_file, std::ios::trunc | std::ios::binary);
  if (!ofs.is_open()) {
    LOG(ERROR) << "open feat_3d_info_file fail: " + feat_3d_info_file;
    return false;
  }

  std::size_t num_points = mFeature_infos.size();
  ofs.write((const char*) &num_points, sizeof(num_points));
  //each 3d points a line
  for (const auto& feat_3d_info : mFeature_infos) {
    ofs.write((const char*) &feat_3d_info.mPoint, sizeof(Point3D));
    std::size_t view_size = feat_3d_info.mView_list.size();
    ofs.write((const char*) &view_size, sizeof(size_t));
    //save view
    for (const auto& view : feat_3d_info.mView_list) {
      ofs.write((const char*) &view, sizeof(View));
    }
    //save decriptor num_desc = feat_3d_info.mDescriptor.size();
    for (const auto& desc : feat_3d_info.mDescriptors) {
      ofs.write((const char*) desc.ptrDesc, desc.desc_legth);
    }
  }
  const bool bOk = ofs.good();
  ofs.close();
  return bOk;
}

bool BundlerParser::LoadFeatures3DInfoBin(const std::string& feat_3d_info_file) {
  std::ifstream ifs(feat_3d_info_file, std::ios::_Nocreate | std::ios::binary);
  if (!ifs.is_open()) {
    LOG(WARNING) << feat_3d_info_file + " file open fail.";
    return false;
  }

  size_t num_points = 0;
  ifs.read((char*) &num_points, sizeof(std::size_t));
 
  VLOG(1) << "load feat_3d_info, total num of points:" << num_points << endl;
  mFeature_infos.clear();
  mFeature_infos.resize(num_points);
  size_t view_list_len = 0;

  for (size_t i = 0; i < num_points; ++i) {
    auto& feat_3d_point = mFeature_infos[i].mPoint;
    ifs.read((char*) &feat_3d_point, sizeof(Point3D));
    ifs.read((char*) &view_list_len, sizeof(size_t));
    //  load view_list
    mFeature_infos[i].mView_list.resize(view_list_len);
    for (auto& view : mFeature_infos[i].mView_list) {
      ifs.read((char*) &view, sizeof(View));
    }

    //load descriptor num_desc = feat_3d_info.mDescriptor.size();
    mFeature_infos[i].mDescriptors.resize(view_list_len);
    for (int j = 0; j < view_list_len; ++j) {
      auto& desc = mFeature_infos[i].mDescriptors[j];
      desc.ptrDesc = new unsigned char[desc.desc_legth];
      if (!desc.ptrDesc) {
        LOG(WARNING) << "new unsigned char fail.try again";
        --j;
        continue;
      }
      ifs.read((char*) desc.ptrDesc, desc.desc_legth);
    }
    VLOG(3) << "load the " << i << "th 3d point okay";
  }

  const bool bOk = !ifs.bad();
  ifs.close();
  LOG(INFO) << "load feat 3d info okay ? " << bOk;
  return bOk;
}

void BundlerParser::GetQueryPictures(std::vector<Picture>& pics) const {
  pics.reserve(mPictures.PictureNum());
  for (const auto& pic : mPictures.Pictures()) {
    if (!pic.Flag()) {
      pics.push_back(pic);
    }
  }
  pics.shrink_to_fit();
  VLOG(1) << "find " << pics.size() << " query pictures";
}

void BundlerParser::GetQueryCameras(std::vector<Camera>& cams) const {
  cams.reserve(mCameras.size());
  const auto& pics = mPictures.Pictures();
  CHECK_EQ(mCameras.size(), pics.size()) << "camera and picture size not equal";
  for (int i = 0; i < mCameras.size(); ++i) {
    if (!pics[i].Flag()) {
      cams.push_back(mCameras[i]);
    }
  }
  cams.shrink_to_fit();
}

void BundlerParser::StatistisDistance(const std::string& file) {
  std::vector<std::vector<int>> distances_vec(mDescriptorMat.rows, std::vector<int>(512, 0));
  std::cout << "StatistisDistance, desc number: " << mDescriptorMat.rows << std::endl;
  std::vector<std::pair<int, int>> same_dists;
  same_dists.reserve(mDescriptorMat.rows);

#pragma omp parallel for schedule(dynamic)
  for (int i = 0; i < mDescriptorMat.rows; ++i) {
    for (int j = i + 1; j < mDescriptorMat.rows; ++j) {
      const auto& desc_i = mDescriptorMat.row(i);
      const auto& desc_j = mDescriptorMat.row(j);
      int dist = cv::norm(desc_i, desc_j, cv::NORM_HAMMING);
      ++distances_vec[i][dist];
      if (dist == 0) {
#pragma omp critical
        same_dists.emplace_back(i, j);
      }
#pragma omp atomic
      ++distances_vec[j][dist];
    }
  }

  std::ofstream ofs{file, std::ios::trunc};
  std::cout << "start write result, rows: " << mDescriptorMat.rows
    << " desc cols: " << mDescriptorMat.cols << std::endl;
  for (int i = 0; i < distances_vec.size(); ++i) {
    for (int k : distances_vec[i]) {
      ofs << k << " ";
    }
    ofs << std::endl;
  }
  ofs.close();
  std::cout << "write finish" << std::endl;

  ofs.open(file + ".sample.txt", std::ios::trunc);
  std::cout << "start write same distance result" << std::endl;
  for (int i = 0; i < 100 && i < same_dists.size(); ++i) {
    const auto& desc_i = mDescriptorMat.row(same_dists[i].first);
    const auto& desc_j = mDescriptorMat.row(same_dists[i].second);
    for (int k = 0; k < desc_i.cols; ++k) {
      ofs << int(desc_i.at<unsigned char>(0, k)) << " ";
    }
    ofs << std::endl;

    for (int k = 0; k < desc_j.cols; ++k) {
      ofs << int(desc_j.at<unsigned char>(0, k)) << " ";
    }
    ofs << std::endl;
  }
  ofs.close();
  std::cout << "write finish" << std::endl;
}


uint64_t BundlerParser::CalculateHammingDistance(const unsigned char* d1,
                                                 const unsigned char* d2,
                                                 int length) {
  uint64_t d1_ = 0;
  uint64_t d2_ = 0;
  uint64_t ret = 0;
  for (int i = 0; i < 8; ++i) {
    std::memcpy(&d1_, d1, 8);
    std::memcpy(&d2_, d2, 8);
    std::advance(d1, 8);
    std::advance(d2, 8);
    ret += __popcnt64(d1_^d2_);
  }
  return ret;
}